import os
import cv2
import argparse
import matplotlib.pyplot as plt
from tqdm import tqdm
from pathlib import Path

def video_to_frames(input_dir, output_dir, visualize=False):
    """
    将input目录中的所有视频文件每秒截取一张图片，并保存到output目录下对应的子目录中
    :param input_dir: 输入视频目录
    :param output_dir: 输出图片目录
    """
    # 确保输入目录存在
    if not os.path.exists(input_dir):
        print(f"错误：输入目录不存在: {input_dir}")
        return
    
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 获取输入目录中的所有视频文件
    video_extensions = ['.mp4', '.avi', '.mov', '.mkv']
    video_files = []
    
    for root, dirs, files in os.walk(input_dir):
        for file in files:
            if any(file.lower().endswith(ext) for ext in video_extensions):
                video_files.append(os.path.join(root, file))
    
    if not video_files:
        print(f"警告：在输入目录 '{input_dir}' 中未找到视频文件")
        return
    
    print(f"找到 {len(video_files)} 个视频文件，开始处理...")
    
    # 存储每个视频的切片数量
    video_frame_counts = []
    video_names = []
    
    # 处理每个视频文件
    for video_path in video_files:
        saved_count = process_video(video_path, output_dir)
        video_name = os.path.splitext(os.path.basename(video_path))[0]
        video_names.append(video_name)
        video_frame_counts.append(saved_count)
    
    print("所有视频处理完成！")
    
    # 生成可视化图表
    if visualize:
        # 设置中文字体
        plt.rcParams["font.sans-serif"] = ["SimHei", "Microsoft YaHei", "Arial Unicode MS"]
        # 解决负号显示问题
        plt.rcParams['axes.unicode_minus'] = False  

        # 可选：设置matplotlib日志级别，忽略字体警告
        import logging
        logging.getLogger('matplotlib.font_manager').setLevel(logging.ERROR)
        plt.figure(figsize=(10, 6))
        plt.bar(video_names, video_frame_counts, color='skyblue')
        plt.xlabel('视频名称')
        plt.ylabel('切片数量')
        plt.title('每个视频的切片数量')
        plt.xticks(rotation=45, ha='right')
        plt.tight_layout()
        
        # 保存图表
        chart_path = os.path.join(output_dir, 'video_frame_counts.png')
        plt.savefig(chart_path)
        print(f"可视化图表已保存至: {chart_path}")
        
        # 显示图表
        plt.show()

def process_video(video_path, output_dir):
    """
    处理单个视频文件，每秒截取一张图片
    :param video_path: 视频文件路径
    :param output_dir: 输出图片目录
    """
    # 获取视频文件名（不含扩展名）
    video_name = os.path.splitext(os.path.basename(video_path))[0]
    
    # 创建输出子目录
    video_output_dir = os.path.join(output_dir, video_name)
    if not os.path.exists(video_output_dir):
        os.makedirs(video_output_dir)
    
    print(f"处理视频: {video_path}")
    print(f"输出目录: {video_output_dir}")
    
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    
    if not cap.isOpened():
        print(f"错误：无法打开视频文件: {video_path}")
        return
    
    # 获取视频帧率
    fps = cap.get(cv2.CAP_PROP_FPS)
    if fps <= 0:
        fps = 30  # 默认为30fps
        print(f"警告：无法获取视频帧率，使用默认值 {fps}fps")
    
    # 获取视频总帧数
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = total_frames / fps
    print(f"视频时长: {duration:.2f}秒，帧率: {fps:.2f}fps")
    
    # 每一帧截取一张图片
    frame_interval = 1  # 每一帧都保存
    current_frame = 0
    saved_frame_count = 0
    
    # 创建进度条
    with tqdm(total=total_frames, unit='帧', desc=f'处理视频: {video_name}', ncols=80) as pbar:
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # 每一帧截取一张图片
            if current_frame % frame_interval == 0:
                # 计算当前时间（秒）
                current_time = current_frame / fps
                
                # 保存图片，文件名包含帧编号和时间
                frame_filename = os.path.join(video_output_dir, f"frame_{current_frame:06d}_{current_time:.2f}s.jpg")
                cv2.imwrite(frame_filename, frame)
                saved_frame_count += 1
            
            current_frame += 1
            pbar.update(1)  # 更新进度条
    
    # 释放资源
    cap.release()
    
    print(f"视频 '{video_name}' 处理完成，共保存 {saved_frame_count} 张图片")
    return saved_frame_count

def main():
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='视频转图片工具 - 每秒截取一张图片')
    parser.add_argument('-i', '--input', default='./input', help='输入视频目录（默认：./input）')
    parser.add_argument('-o', '--output', default='./output', help='输出图片目录（默认：./output）')
    parser.add_argument('-v', '--visualize', action='store_true', help='生成每个视频切片数量的可视化图表')
    
    # 解析命令行参数
    args = parser.parse_args()
    
    # 处理视频文件
    video_to_frames(args.input, args.output, args.visualize)


if __name__ == "__main__":
    main()